#!/usr/bin/python3
##

# filename: csv-measure.py
#
# version: 1.0.0
# author: fangbaoshun
# date: ????-??-??
# description: Calculate launch digest of CSV/CSV2 VMs, including firmware,
#              hashes of initrd+kernel+kernel params.
#
# version: 1.1.0
# author: hanliyang
# date: 2024-12-13
# description: Support calculate launch digest of CSV3 VMs. When calculating
#              the launch digest, you should provide `the number of VM's vcpus`,
#              the `the VM's vcpu model`.
#
#              You can specify the vcpu model by passing parameters in two forms:
#              either `--family $X --model $Y --stepping $Z` or `--vcpumodel
#              $QEMUMODEL`. These two forms of parameters are mutually exclusive.
#
#              If you want launch a kata confidential VM using default VCPUModel
#              (i.e. "host"), you should specify `--family $X --model $Y --stepping
#              $Z`, the $X is physical CPU family number of the machine hosting the
#              kata VM, the $Y is physical CPU model number of the machine hosting
#              the kata VM, and $Z is physical CPU stepping number of the machine
#              hosting the kata VM. You can get $X by `lscpu | grep "^CPU family:"
#              | awk -F':' '{print $2}'` on the target host machine, get $Y by
#              `lscpu | grep "^Model:" | awk -F':' '{print $2}'` on the target host
#              machine, get $Z by `lscpu | grep "^Stepping:" | awk -F':' '{print $2}'`
#              on the target host machine.
#
#              If you want launch a kata confidential VM using the VCPUModel
#              supported by qemu, you should specify `--vcpumodel $QEMUMODEL`. The
#              $QEMUMODEL is the model which supported by qemu. Currently, the qemu
#              support the following models:
#                  "host", "Dhyana", "Dhyana-v1", "Dhyana-v2", "Dhyana-v3",
#                  "Dharma", "Dharma-v1".
#              Be careful, the "host" is a special model, if you execute this script
#              on the machine which hosting the kata confidential VM, you are allowed
#              to specify `--vcpumodel "host"`. Please make sure you know what you
#              are doing when specifying `--vcpumodel "host"`.
#
# version histroy:
# 1.0.0 - ????-??-?? - Support calculate launch digest of CSV/CSV2 VMs.
# 1.1.0 - 2024-12-13 - Support calculate launch digest of CSV3 VMs.

#
## date: 20250304
## Upload this script to gitee.com/hanliyang-kata-coco/deployment.
## We use gmssl instead of pysmx in this script because the pysmx
## in PyPI is not for GuoMi. 
#

import sys
import os
import base64
import hashlib
import subprocess
from gmssl import sm3, func
from argparse import ArgumentParser
from uuid import UUID

#
## GUIDs in ovmf footer table
#
OVMF_TABLE_FOOTER_GUID = "96b582de-1fb2-45f7-baea-a366c55a082d"
CSV3_RESET_BLOCK_GUID = "00f771de-1a7e-4fcb-890e-68c77e2fb44e"
CSV_HASH_TABLE_RV_GUID = "7255371f-3a3b-4b04-927b-1da6efa8d454"

OVMF_FOOTER_TABLE_ENTRY_SIZE = 0x12

#
## GUIDs in hashes table
#
CSV_HASH_TABLE_HEADER_GUID = "9438d606-4f22-4cc9-b479-a793d411fd21"

CSV_KERNEL_ENTRY_GUID = "4de79437-abd2-427f-b835-d5b172d2045b"
CSV_INITRD_ENTRY_GUID = "44baf731-3a2f-4bd7-9af1-41e29169781d"
CSV_CMDLINE_ENTRY_GUID = "97d02dd8-bd20-4c94-aa78-e7714d36ab2a"

HASHES_TABLE_SIZE = 0xB0 # (16 + 2 + 3 * (16 + 2 + 32)) + 8

#
## Definitions of CSV3 measure update extra data
#
CSV3_LAUNCH_ENCRYPT_DATA = 0x01
CSV3_LAUNCH_ENCRYPT_VMCB = 0x02

CSV3_MEASURE_UPDATE_EXTRA_DATA_SIZE = 0x14
CSV3_MEASURE_UPDATE_VCPU_EXTRA_DATA_SIZE = 0x14

#
## VCPU models
#
QEMU_VCPU_MODEL_DICT = {
    'Dhyana': 0x900f01,
    'Dhyana-v1': 0x900f01,
    'Dhyana-v2': 0x900f01,
    'Dhyana-v3': 0x900f01,
    'Dharma': 0x900f40,
    'Dharma-v1': 0x900f40,
}

#
## VMSA
#
VMSA_PAGE_SIZE = 4096


#
## Helper to transform UUID string to UUIDLE bytearray
#
def guid_to_le(guid_str):
    return UUID("{" + guid_str + "}").bytes_le

#
## Construct hashes table
#
def construct_csv_hashes_page(kernel_hash, initrd_hash, cmdline_hash):
    ht_len = 16 + 2 + 3 * (16 + 2 + 32)
    ht_len_aligned = (ht_len + 15) & ~15
    ht = bytearray(ht_len_aligned)

    # Table header
    ht[0:16] = guid_to_le(CSV_HASH_TABLE_HEADER_GUID)
    ht[16:18] = ht_len.to_bytes(2, byteorder='little')

    # Entry 0: kernel command-line
    e = 18
    ht[e:e+16] = guid_to_le(CSV_CMDLINE_ENTRY_GUID)
    ht[e+16:e+18] = (16 + 2 + 32).to_bytes(2, byteorder='little')
    ht[e+18:e+18+32] = cmdline_hash

    # Entry 1: initrd
    e = e+18+32
    ht[e:e+16] = guid_to_le(CSV_INITRD_ENTRY_GUID)
    ht[e+16:e+18] = (16 + 2 + 32).to_bytes(2, byteorder='little')
    ht[e+18:e+18+32] = initrd_hash

    # Entry 2: kernel
    e = e+18+32
    ht[e:e+16] = guid_to_le(CSV_KERNEL_ENTRY_GUID)
    ht[e+16:e+18] = (16 + 2 + 32).to_bytes(2, byteorder='little')
    ht[e+18:e+18+32] = kernel_hash

    return ht

#
## Parse footer table from the ovmf file
#
def parse_ovmf_footer_table(firmware_path):
    ovmf_footer_table_dict = {}

    with open(firmware_path, 'rb') as fh:
        buf = fh.read()

        # check the ovmf table footer
        footer_start = len(buf) - 32 - OVMF_FOOTER_TABLE_ENTRY_SIZE
        footer_buf = buf[footer_start:]
        if footer_buf[2:OVMF_FOOTER_TABLE_ENTRY_SIZE] != guid_to_le(OVMF_TABLE_FOOTER_GUID):
            print("Error: not found ovmf table footer in", firmware_path)
            sys.exit(1)

        # get size of ovmf footer table from the ovmf table footer
        table_size = int.from_bytes(footer_buf[:2], byteorder='little', signed=False) - OVMF_FOOTER_TABLE_ENTRY_SIZE
        table_buf = buf[footer_start - table_size:footer_start]
        # traverse whole ovmf footer table, save the table entries
        while len(table_buf) >= OVMF_FOOTER_TABLE_ENTRY_SIZE:
            table_size = len(table_buf)

            entry_hdr_buf = table_buf[table_size - OVMF_FOOTER_TABLE_ENTRY_SIZE:]
            entry_size = int.from_bytes(entry_hdr_buf[:2], byteorder='little', signed=False)
            if entry_size < OVMF_FOOTER_TABLE_ENTRY_SIZE:
                print("Error: invalid size of ovmf footer entry")
                sys.exit(1)
            entry_data_buf = table_buf[table_size - entry_size:table_size - OVMF_FOOTER_TABLE_ENTRY_SIZE]

            # append entry to dict
            ovmf_footer_table_dict[entry_hdr_buf[2:OVMF_FOOTER_TABLE_ENTRY_SIZE]] = entry_data_buf

            table_buf = table_buf[:table_size - entry_size]

    return ovmf_footer_table_dict

#
## Construct CPU's Family, Model, Stepping Identifiers
#
def construct_vcpu_sig(args):
    if args.vcpumodel is not None:
        if args.vcpumodel != "host":
            return QEMU_VCPU_MODEL_DICT[args.vcpumodel]
        else:
            shell_command = 'lscpu | grep \"^CPU family:\" | awk -F\':\' \'{print $2}\''
            result = subprocess.run(shell_command, shell=True, capture_output=True, text=True)
            if result.returncode == 0:
                cpu_family = int(result.stdout.strip())
            shell_command = 'lscpu | grep \"^Model:\" | awk -F\':\' \'{print $2}\''
            result = subprocess.run(shell_command, shell=True, capture_output=True, text=True)
            if result.returncode == 0:
                cpu_model = int(result.stdout.strip())
            shell_command = 'lscpu | grep \"^Stepping:\" | awk -F\':\' \'{print $2}\''
            result = subprocess.run(shell_command, shell=True, capture_output=True, text=True)
            if result.returncode == 0:
                cpu_stepping = int(result.stdout.strip())
    else:
        cpu_family = args.family
        cpu_model = args.model
        cpu_stepping = args.stepping

    # See CPUID_Fn0000_0001_EAX Family, Model, Stepping Identifiers
    if cpu_family > 0xF:
        base_family = 0xF
        ext_family = (cpu_family - 0x0F) & 0xFF
    else:
        base_family = cpu_family
        ext_family = 0
    base_model = cpu_model & 0xF
    ext_model = (cpu_model >> 4) & 0xF
    stepping = cpu_stepping & 0xF

    return (ext_family << 20) | (ext_model << 16) | (base_family << 8) | (base_model << 4) | stepping

#
## Helper to construct VMSA page
#
def construct_vmsa_page(vcpu_id, reset_eip, vcpu_sig):
    vmsa_page = bytearray(bytes(VMSA_PAGE_SIZE))

    # ES, CS, SS, DS, FS, GS
    offset = 0x0
    segreg_selector = 0x0
    segreg_attrib = 0x92 | 0x1
    segreg_limit = 0xFFFF
    segreg_base = 0x0
    vmsa_page[offset:offset + 2] = segreg_selector.to_bytes(2, byteorder='little')
    vmsa_page[offset + 2:offset + 4] = segreg_attrib.to_bytes(2, byteorder='little')
    vmsa_page[offset + 4:offset + 8] = segreg_limit.to_bytes(4, byteorder='little')
    vmsa_page[offset + 8:offset + 16] = segreg_base.to_bytes(8, byteorder='little')
    offset = 0x10
    segreg_selector = 0xF000
    segreg_attrib = 0x9A | 0x1
    segreg_limit = 0xFFFF
    if vcpu_id == 0:
        segreg_base = 0xFFFF0000
    else:
        segreg_base = reset_eip & 0xFFFF0000
    vmsa_page[offset:offset + 2] = segreg_selector.to_bytes(2, byteorder='little')
    vmsa_page[offset + 2:offset + 4] = segreg_attrib.to_bytes(2, byteorder='little')
    vmsa_page[offset + 4:offset + 8] = segreg_limit.to_bytes(4, byteorder='little')
    vmsa_page[offset + 8:offset + 16] = segreg_base.to_bytes(8, byteorder='little')
    offset = 0x20
    segreg_selector = 0x0
    segreg_attrib = 0x92 | 0x1
    segreg_limit = 0xFFFF
    segreg_base = 0x0
    vmsa_page[offset:offset + 2] = segreg_selector.to_bytes(2, byteorder='little')
    vmsa_page[offset + 2:offset + 4] = segreg_attrib.to_bytes(2, byteorder='little')
    vmsa_page[offset + 4:offset + 8] = segreg_limit.to_bytes(4, byteorder='little')
    vmsa_page[offset + 8:offset + 16] = segreg_base.to_bytes(8, byteorder='little')
    offset = 0x30
    vmsa_page[offset:offset + 2] = segreg_selector.to_bytes(2, byteorder='little')
    vmsa_page[offset + 2:offset + 4] = segreg_attrib.to_bytes(2, byteorder='little')
    vmsa_page[offset + 4:offset + 8] = segreg_limit.to_bytes(4, byteorder='little')
    vmsa_page[offset + 8:offset + 16] = segreg_base.to_bytes(8, byteorder='little')
    offset = 0x40
    vmsa_page[offset:offset + 2] = segreg_selector.to_bytes(2, byteorder='little')
    vmsa_page[offset + 2:offset + 4] = segreg_attrib.to_bytes(2, byteorder='little')
    vmsa_page[offset + 4:offset + 8] = segreg_limit.to_bytes(4, byteorder='little')
    vmsa_page[offset + 8:offset + 16] = segreg_base.to_bytes(8, byteorder='little')
    offset = 0x50
    vmsa_page[offset:offset + 2] = segreg_selector.to_bytes(2, byteorder='little')
    vmsa_page[offset + 2:offset + 4] = segreg_attrib.to_bytes(2, byteorder='little')
    vmsa_page[offset + 4:offset + 8] = segreg_limit.to_bytes(4, byteorder='little')
    vmsa_page[offset + 8:offset + 16] = segreg_base.to_bytes(8, byteorder='little')
    # GDTR, LDTR, IDTR, TR
    offset = 0x60 + 2 + 2
    segreg_limit = 0xFFFF
    vmsa_page[offset:offset + 4] = segreg_limit.to_bytes(4, byteorder='little')
    offset = 0x70 + 2
    segreg_attrib = 0x82
    segreg_limit = 0xFFFF
    vmsa_page[offset:offset + 2] = segreg_attrib.to_bytes(2, byteorder='little')
    vmsa_page[offset + 2:offset + 6] = segreg_limit.to_bytes(4, byteorder='little')
    offset = 0x80 + 2 + 2
    segreg_limit = 0xFFFF
    vmsa_page[offset:offset + 4] = segreg_limit.to_bytes(4, byteorder='little')
    offset = 0x90 + 2
    segreg_attrib = 0x83 | 0x8
    segreg_limit = 0xFFFF
    vmsa_page[offset:offset + 2] = segreg_attrib.to_bytes(2, byteorder='little')
    vmsa_page[offset + 2:offset + 6] = segreg_limit.to_bytes(4, byteorder='little')
    # PL0_SSP, PL1_SSP, PL2_SSP, PL3_SSP
    # U_CET, RESERVED, VMPL
    # CPL (== 0)
    # RESERVED
    # EFER
    offset = 0xD0
    efer = 0x0 | (1 << 12)
    vmsa_page[offset:offset + 8] = efer.to_bytes(8, byteorder='little')
    # RESERVED
    # XSS (== 0)
    # CR4, CR3, CR0
    offset = 0x148
    cr4 = 0x0 | (1 << 6)
    vmsa_page[offset:offset + 8] = cr4.to_bytes(8, byteorder='little')
    offset = 0x158
    cr0 = 0x0000000060000010 & ~(0x3 << 29)
    vmsa_page[offset:offset + 8] = cr0.to_bytes(8, byteorder='little')
    # DR7, DR6
    offset = 0x160
    dr7 = 0x0000000000000400
    vmsa_page[offset:offset + 8] = dr7.to_bytes(8, byteorder='little')
    offset = 0x168
    dr6 = 0x00000000FFFF0FF0
    vmsa_page[offset:offset + 8] = dr6.to_bytes(8, byteorder='little')
    # RFLAGS
    offset = 0x170
    rflags = 0x0000000000000002
    vmsa_page[offset:offset + 8] = rflags.to_bytes(8, byteorder='little')
    # RIP
    offset = 0x178
    if vcpu_id == 0:
        rip = 0x000000000000FFF0
    else:
        rip = reset_eip & 0x0000FFFF
    vmsa_page[offset:offset + 8] = rip.to_bytes(8, byteorder='little')
    # DR0, DR1, DR2, DR3
    # DR0_ADDR_MASK, DR1_ADDR_MASK, DR2_ADDR_MASK, DR3_ADDR_MASK
    # RESERVED
    # RSP (== 0)
    # S_CET, SSP, ISST_ADDR
    # RAX (== 0)
    # STAR (== 0), LSTAR (== 0), CSTAR (== 0), SFMASK (== 0), KERNLE_GS_BASE (== 0)
    # SYSENTER_CS (== 0), SYSENTER_ESP (== 0), SYSENTER_EIP (== 0)
    # CR2 (== 0)
    # RESERVED
    # G_PAT
    offset = 0x268
    g_pat = 0x0007040600070406
    vmsa_page[offset:offset + 8] = g_pat.to_bytes(8, byteorder='little')
    # DBGCTL (== 0), BR_FROM (== 0), BR_TO (== 0), LASTEXCPFROM (== 0), LASTEXCPTO (== 0)
    # RESERVED
    # PKRU, TSC_AUX
    # RESERVED
    # RCX (== 0)
    # RDX
    offset = 0x310
    vmsa_page[offset:offset + 8] = vcpu_sig.to_bytes(8, byteorder='little')
    # RBX (== 0)
    # RESERVED
    # RBP, RSI, RDI, R8, R9, R10, R11, R12, R13, R14, R15
    # RESERVED
    # GUEST_EXITINFO1, GUEST_EXITINFO2, GUEST_EXITINTINFO, GUEST_NRIP
    # xxx_FEATURES, VINTR_CTRL
    # GUEST_EXITCODE, VIRTUAL_TOM
    # TLB_ID, PCPU_ID
    # EVENTINJ
    # XCR0
    offset = 0x3E8
    xcr0 = 0x0000000000000001
    vmsa_page[offset:offset + 8] = xcr0.to_bytes(8, byteorder='little')
    # RESERVED
    # X87_DP, MXCSR, X87_FTW, X87_FSW
    # X87_FCW, X87_FOP, X87_DS, X87_CS, X87_RIP
    # FPREG_X87, FPREG_XMM, FPREG_YMM
    return vmsa_page

#
## Helper to construct CSV3 measure update extra data
#
def construct_csv3_measure_update_extra_data(base, size):
    data = bytearray(bytes(CSV3_MEASURE_UPDATE_EXTRA_DATA_SIZE))

    # fill in update ID
    data[0:4] = CSV3_LAUNCH_ENCRYPT_DATA.to_bytes(4, byteorder='little')
    data[4:12] = base.to_bytes(8, byteorder='little')
    data[12:20] = size.to_bytes(8, byteorder='little')

    return data

#
## Helper to construct CSV3 measure update_vcpu extra data
#
def construct_csv3_measure_update_vcpu_extra_data(vcpu_id):
    data = bytearray(bytes(CSV3_MEASURE_UPDATE_VCPU_EXTRA_DATA_SIZE))

    # fill in update ID
    data[0:4] = CSV3_LAUNCH_ENCRYPT_VMCB.to_bytes(4, byteorder='little')
    data[4:12] = vcpu_id.to_bytes(8, byteorder='little')
    data[12:20] = VMSA_PAGE_SIZE.to_bytes(8, byteorder='little')

    return data

#
## Check whether the combination of parameters is correct
#
def check_arguments(args):
    if args.csv3 is True:
        if args.smp is None:
            print("Error: missing '--smp X'")
            sys.exit(1)
        if args.vcpumodel is None:
            if args.family is None or args.model is None or args.stepping is None:
                print("Error: missing '--vcpumodel \"qemumodel\"' or '--family X --model Y --stepping Z'")
                sys.exit(1)
        if args.vcpumodel is not None:
            if args.family is not None or args.model is not None or args.stepping is not None:
                print("Error: '--vcpumodel \"qemumodel\"' and '--family X --model Y --stepping Z' is mutually exclusive")
                sys.exit(1)

            vcpu_sig = None
            if args.vcpumodel != "host":
                try:
                    vcpu_sig = QEMU_VCPU_MODEL_DICT[args.vcpumodel]
                except Exception as e:
                    print(f"Error: vcpumodel is invalid: {e}")
                    sys.exit(1)

class SM3Wrapper:
    def __init__(self):
        self.data_list = []

    def update(self, data):
        # 将输入的数据（字节类型）转换为列表并添加到存储数据的列表中
        data_list = func.bytes_to_list(data)
        self.data_list.extend(data_list)

    def digest(self):
        # 调用 gmssl 库中的 sm3_hash 函数计算最终的哈希值
        hash_result = sm3.sm3_hash(self.data_list)
        return bytes.fromhex(hash_result)

def main(args):
    kernel_hash = None
    initrd_hash = None
    cmdline_hash = None
    csv_hashes_table = None

    ovmf_footer_table_dict = {}

    if args.kernel:
        with open(args.kernel, 'rb') as fh:
            h = hashlib.sha256(fh.read())
            kernel_hash = h.digest()

    if args.initrd:
        with open(args.initrd, 'rb') as fh:
            h = hashlib.sha256(fh.read())
            initrd_hash = h.digest()

    if args.cmdline:
        with open(args.cmdline, 'rb') as fh:
            cmd_buf = fh.read()
            cmdline_buffer = cmd_buf[:-1]
            cmdline_buffer = cmdline_buffer + b'\x00'

            h = hashlib.sha256(cmdline_buffer)
            cmdline_hash = h.digest()

    if kernel_hash and initrd_hash and cmdline_hash:
        csv_hashes_table = construct_csv_hashes_page(kernel_hash, initrd_hash, cmdline_hash)

    # parse ovmf footer table for CSV3
    if args.csv3 is True:
        ovmf_footer_table_dict = parse_ovmf_footer_table(args.ovmf)

    with open(args.ovmf, 'rb') as fh:
        sm3_hash = SM3Wrapper()

        # Update firmware
        if args.csv3 is True:
            # construct measure update extra data
            ovmf_size = os.path.getsize(args.ovmf)
            ovmf_base = 0x100000000 - ovmf_size
            update_extra_data = construct_csv3_measure_update_extra_data(ovmf_base, ovmf_size)
            sm3_hash.update(update_extra_data)

        sm3_hash.update(fh.read())

        if csv_hashes_table is not None:
            if args.csv3 is True:
                # construct measure update extra data
                entry_data = ovmf_footer_table_dict[guid_to_le(CSV_HASH_TABLE_RV_GUID)]
                ht_base = int.from_bytes(entry_data[:4], byteorder='little', signed=False)
                if ht_base == 0 and int.from_bytes(entry_data[4:8], byteorder='little', signed=False) == 0:
                    print("Error: hashes table area is invalid")
                    sys.exit(1)
                ht_size = HASHES_TABLE_SIZE
                update_extra_data = construct_csv3_measure_update_extra_data(ht_base, ht_size)
                sm3_hash.update(update_extra_data)

            sm3_hash.update(csv_hashes_table)

        # Update VMSA, if CSV3
        if args.csv3 is True:
            entry_data = ovmf_footer_table_dict[guid_to_le(CSV3_RESET_BLOCK_GUID)]
            reset_eip = int.from_bytes(entry_data[:4], byteorder='little', signed=False)
            vcpu_sig = construct_vcpu_sig(args)

            i = 0
            while i < args.smp:
                # construct measure update VCPU extra data
                update_vcpu_extra_data = construct_csv3_measure_update_vcpu_extra_data(i)
                sm3_hash.update(update_vcpu_extra_data)
                # construct VMSA
                vmsa_page = construct_vmsa_page(i, reset_eip, vcpu_sig)
                sm3_hash.update(vmsa_page)
                i = i + 1

        final_hash = sm3_hash.digest()
        final_hash_str = str(base64.b64encode(final_hash), 'utf-8')
        print(final_hash_str)

if __name__ == "__main__":
    parser = ArgumentParser(description='Calculate firmware digest')

    parser.add_argument('--ovmf',
                        help='location of OVMF file to calculate hash from',
                        required=True)
    parser.add_argument('--kernel',
                        help='location of kernel file to calculate hash from')
    parser.add_argument('--initrd',
                        help='location of initrd file to calculate hash from')
    parser.add_argument('--cmdline',
                        help='the kernel command line to calculate hash from')
    parser.add_argument('--csv3', action='store_true',
                        help='the calculation is for CSV3 VM')
    parser.add_argument('--smp', type=int,
                        help='the number of vcpus for the VM, the value is of type int')
    parser.add_argument('--family', type=int,
                        help='the family number of VM\'s vcpu, the value is of type int')
    parser.add_argument('--model', type=int,
                        help='the model number of VM\'s vcpu, the value is of type int')
    parser.add_argument('--stepping', type=int,
                        help='the stepping number of VM\'s vcpu, the value is of type int')
    parser.add_argument('--vcpumodel', type=str,
                        help='the stepping number of VM\'s vcpu, the value is of type string.\n'
                             'This argument is mutually exclusive with `--famliy $X --model $Y --stepping $Z`.\n'
                             'The valid value of this argument are:\n'
                             '   "host", "Dhyana", "Dhyana-v1", "Dhyana-v2", "Dhyana-v3"\n'
                             '   "Dharma", "Dharma-v1"')

    args = parser.parse_args()

    # validate the arguments
    check_arguments(args)

    main(args)
